import google.generativeai as genai
from google.generativeai import types
from PIL import Image
from io import BytesIO
import base64
from typing import Dict, List, Optional, Tuple, Union
import logging

class GeminiClient:
    """Gemini API客户端"""
    
    def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash-preview-image-generation"):
        """
        初始化Gemini客户端
        
        Args:
            api_key: Gemini API密钥
            model_name: 模型名称
        """
        if not api_key:
            raise ValueError("Gemini API密钥不能为空")
        
        self.api_key = api_key
        self.model_name = model_name
        
        # 配置API
        genai.configure(api_key=api_key)
        
        try:
            self.client = genai.Client()
            # 测试连接
            self._test_connection()
        except Exception as e:
            logging.error(f"初始化Gemini客户端失败: {e}")
            raise
    
    def _test_connection(self):
        """测试API连接"""
        try:
            # 简单的测试请求
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=["测试连接"],
                config=types.GenerateContentConfig(
                    response_modalities=['TEXT']
                )
            )
            logging.info("Gemini API连接测试成功")
        except Exception as e:
            logging.error(f"Gemini API连接测试失败: {e}")
            raise
    
    def generate_image_from_text(self, prompt: str, conversation_history: List[Dict] = None) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        根据文本提示生成图片
        
        Args:
            prompt: 文本提示
            conversation_history: 对话历史
            
        Returns:
            (生成的图片base64, 响应文本, 错误信息)
        """
        try:
            # 构建内容
            contents = []
            
            # 添加对话历史上下文
            if conversation_history:
                for entry in conversation_history[-3:]:  # 只取最近3轮对话
                    if entry.get('user_input'):
                        contents.append(entry['user_input'])
                    if entry.get('ai_response'):
                        contents.append(entry['ai_response'])
            
            # 添加当前提示
            contents.append(prompt)
            
            # 调用API
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=['TEXT', 'IMAGE']
                )
            )
            
            return self._process_response(response)
            
        except Exception as e:
            error_msg = f"生成图片失败: {str(e)}"
            logging.error(error_msg)
            return None, None, error_msg
    
    def edit_image_with_text(self, image_path: str, prompt: str, conversation_history: List[Dict] = None) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        根据文本提示编辑图片
        
        Args:
            image_path: 输入图片路径
            prompt: 编辑提示
            conversation_history: 对话历史
            
        Returns:
            (生成的图片base64, 响应文本, 错误信息)
        """
        try:
            # 加载图片
            image = Image.open(image_path)
            
            # 构建内容
            contents = []
            
            # 添加对话历史上下文
            if conversation_history:
                for entry in conversation_history[-3:]:  # 只取最近3轮对话
                    if entry.get('user_input'):
                        contents.append(entry['user_input'])
                    if entry.get('ai_response'):
                        contents.append(entry['ai_response'])
            
            # 添加当前提示和图片
            contents.extend([prompt, image])
            
            # 调用API
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=['TEXT', 'IMAGE']
                )
            )
            
            return self._process_response(response)
            
        except Exception as e:
            error_msg = f"编辑图片失败: {str(e)}"
            logging.error(error_msg)
            return None, None, error_msg
    
    def edit_image_with_base64(self, base64_image: str, prompt: str, conversation_history: List[Dict] = None) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        根据base64图片和文本提示编辑图片
        
        Args:
            base64_image: base64编码的图片
            prompt: 编辑提示
            conversation_history: 对话历史
            
        Returns:
            (生成的图片base64, 响应文本, 错误信息)
        """
        try:
            # 处理base64图片
            if base64_image.startswith('data:image/'):
                base64_image = base64_image.split(',', 1)[1]
            
            # 解码图片
            image_bytes = base64.b64decode(base64_image)
            image = Image.open(BytesIO(image_bytes))
            
            # 构建内容
            contents = []
            
            # 添加对话历史上下文
            if conversation_history:
                for entry in conversation_history[-3:]:  # 只取最近3轮对话
                    if entry.get('user_input'):
                        contents.append(entry['user_input'])
                    if entry.get('ai_response'):
                        contents.append(entry['ai_response'])
            
            # 添加当前提示和图片
            contents.extend([prompt, image])
            
            # 调用API
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=['TEXT', 'IMAGE']
                )
            )
            
            return self._process_response(response)
            
        except Exception as e:
            error_msg = f"编辑图片失败: {str(e)}"
            logging.error(error_msg)
            return None, None, error_msg
    
    def _process_response(self, response) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        处理API响应
        
        Args:
            response: Gemini API响应
            
        Returns:
            (生成的图片base64, 响应文本, 错误信息)
        """
        try:
            if not response.candidates:
                return None, None, "API返回空响应"
            
            candidate = response.candidates[0]
            if not candidate.content or not candidate.content.parts:
                return None, None, "响应中没有内容"
            
            text_response = None
            image_base64 = None
            
            # 处理响应部分
            for part in candidate.content.parts:
                if part.text is not None:
                    text_response = part.text
                elif part.inline_data is not None:
                    # 将图片数据转换为base64
                    image_data = part.inline_data.data
                    mime_type = part.inline_data.mime_type
                    image_base64 = f"data:{mime_type};base64,{base64.b64encode(image_data).decode('utf-8')}"
            
            if not text_response and not image_base64:
                return None, None, "响应中没有有效内容"
            
            return image_base64, text_response, None
            
        except Exception as e:
            error_msg = f"处理响应失败: {str(e)}"
            logging.error(error_msg)
            return None, None, error_msg
